import tensorflow as tf
import numpy as np
import cv2
import os
from datetime import datetime
import shutil

class DataAugmentor:
    def __init__(self, target_dir, background_dir, output_dir='dataset_augmented'):
        self.target_dir = target_dir
        self.background_dir = background_dir
        self.output_dir = output_dir
        self.target_output_dir = os.path.join(output_dir, 'target')
        self.background_output_dir = os.path.join(output_dir, 'background')
        
        # 定义数据增强器
        self.augmenter = tf.keras.Sequential([
            tf.keras.layers.RandomFlip("horizontal_and_vertical"),
            tf.keras.layers.RandomRotation(0.3),
            tf.keras.layers.RandomZoom(0.2),
            tf.keras.layers.RandomBrightness(0.2),
            tf.keras.layers.RandomContrast(0.2),
            tf.keras.layers.GaussianNoise(0.1),
        ])
    
    def get_existing_augmented_files(self, output_dir):
        """获取已增强的文件基础名列表"""
        if not os.path.exists(output_dir):
            return set()
        
        existing_files = set()
        for filename in os.listdir(output_dir):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                # 提取基础文件名（去掉 _original 和 _aug_x 后缀）
                base_name = filename.split('_original')[0].split('_aug_')[0]
                existing_files.add(base_name)
        
        return existing_files
    
    def check_augmented_data(self):
        """检查已增强的数据状态"""
        if not os.path.exists(self.output_dir):
            return False, 0, 0, set(), set()
        
        target_files = self.get_existing_augmented_files(self.target_output_dir)
        background_files = self.get_existing_augmented_files(self.background_output_dir)
        
        target_count = len([f for f in os.listdir(self.target_output_dir) 
                          if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
        background_count = len([f for f in os.listdir(self.background_output_dir) 
                              if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
        
        return True, target_count, background_count, target_files, background_files
    
    def augment_images(self, multiplier=3):
        """增强图片数据"""
        print("\n开始数据增强...")
        
        # 检查已增强的数据
        exists, target_count, background_count, existing_targets, existing_backgrounds = self.check_augmented_data()
        
        if exists:
            print(f"\n检测到已存在的增强数据:")
            print(f"- 目标物品图片: {target_count} 张")
            print(f"- 背景图片: {background_count} 张")
            
            while True:
                print("\n请选择处理方式：")
                print("1. 使用现有数据")
                print("2. 仅增强新数据")
                print("3. 重新生成所有数据")
                choice = input("\n请输入选择 (1-3): ")
                
                if choice == '1':
                    print("\n使用现有增强数据")
                    return
                elif choice == '2':
                    print("\n仅增强新数据...")
                    break
                elif choice == '3':
                    print("\n重新生成所有数据...")
                    existing_targets.clear()
                    existing_backgrounds.clear()
                    if os.path.exists(self.output_dir):
                        shutil.rmtree(self.output_dir)
                    break
                else:
                    print("无效输入，请重试")
        
        # 创建输出目录
        os.makedirs(self.target_output_dir, exist_ok=True)
        os.makedirs(self.background_output_dir, exist_ok=True)
        
        # 增强目标物品图片
        self._augment_folder(
            self.target_dir, 
            self.target_output_dir, 
            multiplier, 
            existing_targets
        )
        
        # 增强背景图片
        self._augment_folder(
            self.background_dir, 
            self.background_output_dir, 
            multiplier, 
            existing_backgrounds
        )
        
        print("\n数据增强完成！")
        
        # 显示最终统计信息
        exists, target_count, background_count, _, _ = self.check_augmented_data()
        print(f"\n增强后的数据统计:")
        print(f"- 目标物品图片: {target_count} 张")
        print(f"- 背景图片: {background_count} 张")
    
    def _augment_folder(self, input_dir, output_dir, multiplier, existing_files):
        """对指定文件夹中的图片进行增强"""
        images = [f for f in os.listdir(input_dir) 
                 if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        
        # 过滤出需要增强的图片
        new_images = []
        skipped_images = []
        for img_name in images:
            base_name = os.path.splitext(img_name)[0]
            if base_name in existing_files:
                skipped_images.append(img_name)
            else:
                new_images.append(img_name)
        
        print(f"\n处理文件夹: {input_dir}")
        print(f"- 原始图片总数: {len(images)}")
        print(f"- 跳过已增强: {len(skipped_images)}")
        print(f"- 待增强图片: {len(new_images)}")
        
        for img_name in new_images:
            # 读取原始图片
            img_path = os.path.join(input_dir, img_name)
            img = cv2.imread(img_path)
            if img is None:
                print(f"无法读取图片: {img_path}")
                continue
            
            # 转换颜色空间
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # 复制原始图片到输出目录
            base_name = os.path.splitext(img_name)[0]
            cv2.imwrite(
                os.path.join(output_dir, f"{base_name}_original.jpg"),
                cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            )
            
            # 生成增强图片
            for i in range(multiplier):
                # 转换为张量
                img_tensor = tf.convert_to_tensor(img, dtype=tf.float32)
                img_tensor = tf.expand_dims(img_tensor, 0)  # 添加批次维度
                
                # 应用增强
                augmented = self.augmenter(img_tensor)
                
                # 转换回numpy数组
                augmented = augmented[0].numpy().astype(np.uint8)
                
                # 应用高级增强
                augmented = self.apply_advanced_augmentation(augmented)
                
                # 保存增强后的图片
                output_path = os.path.join(output_dir, f"{base_name}_aug_{i+1}.jpg")
                cv2.imwrite(output_path, cv2.cvtColor(augmented, cv2.COLOR_RGB2BGR))
        
        # 统计增强后的图片数量
        final_count = len([f for f in os.listdir(output_dir) 
                          if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
        print(f"增强后图片数量: {final_count}")
    
    def apply_advanced_augmentation(self, img):
        """应用高级增强技术"""
        augmented = img.copy()
        
        # 随机应用以下增强方法
        if np.random.random() < 0.3:
            # 添加高斯模糊
            kernel_size = np.random.choice([3, 5, 7])
            augmented = cv2.GaussianBlur(augmented, (kernel_size, kernel_size), 0)
            
        if np.random.random() < 0.3:
            # 添加运动模糊
            kernel_size = np.random.choice([3, 5, 7])
            kernel = np.zeros((kernel_size, kernel_size))
            kernel[int((kernel_size-1)/2), :] = np.ones(kernel_size)
            kernel = kernel / kernel_size
            augmented = cv2.filter2D(augmented, -1, kernel)
            
        if np.random.random() < 0.3:
            # 调整饱和度
            hsv = cv2.cvtColor(augmented, cv2.COLOR_RGB2HSV)
            hsv[:, :, 1] = hsv[:, :, 1] * np.random.uniform(0.7, 1.3)
            augmented = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
            
        if np.random.random() < 0.3:
            # 添加椒盐噪声
            prob = np.random.uniform(0.01, 0.05)
            rnd = np.random.random(augmented.shape[:2])
            augmented[rnd < prob] = 0
            augmented[rnd > 1 - prob] = 255
        
        return augmented

def main():
    """主函数"""
    augmentor = DataAugmentor(
        target_dir='dataset/target',
        background_dir='dataset/background'
    )
    augmentor.augment_images(multiplier=3)

if __name__ == "__main__":
    main() 